/**
 * Copyright 2024-2024 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "model_builder/hiai_ir_build_v2.h"

#include <dlfcn.h>

#include "stub_hiai_ir_build_model.h"
#include "framework/infra/log/log_fmk_interface.h"
#include "framework/infra/log/log.h"
#include "infra/base/assertion.h"
#include "framework/graph/utils/graph_utils.h"
#include "framework/graph/utils/attr_utils.h"
#include "framework/compatible/ir_transformer.h"
#include "model/built_model_aipp.h"
namespace hiai {
namespace {
class ExtendedCompatibleModel {
public:
    explicit ExtendedCompatibleModel(ge::Model* model)
    {
        model_ = model;
    }

    ge::Model* model_;
    std::string aippConfig_;
};

bool IsSupportNewAipp()
{
    return false;
}

static hiai::Status MakeAippCompatible(ExtendedCompatibleModel& extendedModel)
{
    if (IsSupportNewAipp()) {
        return hiai::SUCCESS;
    }

    ge::ComputeGraphPtr graph = ge::GraphUtils::GetComputeGraph(extendedModel.model_->GetGraph());
    if (graph == nullptr) {
        FMK_LOGE("ComputeGraphPtr is nullptr");
        return hiai::FAILURE;
    }

    const char* HIAI_IR_BUILD_AIPP = "libhiai_ir_build_aipp_ddk.so";
    void* handle = dlopen(HIAI_IR_BUILD_AIPP, RTLD_LAZY);
    if (handle == nullptr) {
        FMK_LOGE("Create DynamicLoadHelper fail!");
        return hiai::FAILURE;
    }
    auto func = reinterpret_cast<Status (*)(ge::ComputeGraph&, std::string&)>(
        dlsym(handle, "GenerateAippCompatibleInfoAdapter"));
    if (func == nullptr) {
        FMK_LOGE("func is nullptr");
        return hiai::FAILURE;
    }
    return func(*graph, extendedModel.aippConfig_);
}

static std::unique_ptr<ExtendedCompatibleModel> MakeCompatibleModel(const std::shared_ptr<ge::Model>& model)
{
    (void)ge::AttrUtils::SetInt(&*model, "stream_num", 1);
    std::unique_ptr<ExtendedCompatibleModel> compatibleModel = ge::make_unique<ExtendedCompatibleModel>(model.get());

    ge::ComputeGraphPtr graph = ge::GraphUtils::GetComputeGraph(model->GetGraph());
    if (graph == nullptr) {
        FMK_LOGE("Graph is null.");
        return nullptr;
    }

    if (!VerifyIRAPI(graph)) {
        return nullptr;
    }

    if (MakeAippCompatible(*compatibleModel.get()) != hiai::SUCCESS) {
        FMK_LOGE("MakeAippCompatible failed.");
        return nullptr;
    }

    return compatibleModel;
}
} // namespace

static std::shared_ptr<hiai::IBuiltModel> BuildCompatibleModel(const hiai::ModelBuildOptions& options,
    const std::string& modelName, std::unique_ptr<ExtendedCompatibleModel>& compatibleModel)
{
    auto builtModel = BuildModel(options, modelName, *(compatibleModel->model_));
    HIAI_EXPECT_NOT_NULL_R(builtModel, nullptr);

    if (compatibleModel->aippConfig_.size() > 0) {
        hiai::CustomModelData customModelData {hiai::AIPP_PREPROCESS_TYPE, compatibleModel->aippConfig_};
        builtModel->SetCustomData(customModelData);
    }

    FMK_LOGI("Build ir model success.");
    return builtModel;
}

GRAPH_API_EXPORT Status HiaiIrBuildV2::Build(const hiai::ModelBuildOptions& options, const std::string& modelName,
    const std::shared_ptr<ge::Model>& model, std::shared_ptr<hiai::IBuiltModel>& builtModel)
{
    auto compatibleModel = MakeCompatibleModel(model);
    if (compatibleModel == nullptr) {
        FMK_LOGE("MakeCompatibleModel falied");
        return hiai::FAILURE;
    }
    builtModel = BuildCompatibleModel(options, modelName, compatibleModel);
    if (builtModel == nullptr) {
        FMK_LOGE("BuildCompatibleModel falied");
        return hiai::FAILURE;
    }
    return hiai::SUCCESS;
}
}